import numpy as np
from ultralytics import YOLO
import cv2
import cvzone
import math
import torch

#导入排序器
from sort import *

# 检查是否有可用的CUDA设备
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"正在使用的设备: {device}")

# 检查YOLO是否使用GPU
print(f"YOLO是否使用GPU: {device.type == 'cuda'}")

# 如果使用GPU，打印GPU信息
if device.type == 'cuda':
    print(f"GPU名称: {torch.cuda.get_device_name(0)}")
    print(f"GPU内存: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.1f} GB")





#用于视频中
cap=cv2.VideoCapture("../Video/car-1.mp4")


#导入模型
model=YOLO("../Yolo-Weights/yolov8n.pt")


#导入数据集
classNames = [
    "person", "bicycle", "car", "motorbike", "aeroplane", "bus", "train", "truck", "boat",
    "traffic light", "fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat",
    "dog", "horse", "sheep", "cow", "elephant", "bear", "zebra", "giraffe", "backpack", "umbrella",
    "handbag", "tie", "suitcase", "frisbee", "skis", "snowboard", "sports ball", "kite", "baseball bat",
    "baseball glove", "skateboard", "surfboard", "tennis racket", "bottle", "wine glass", "cup",
    "fork", "knife", "spoon", "bowl", "banana", "apple", "sandwich", "orange", "broccoli",
    "carrot", "hot dog", "pizza", "donut", "cake", "chair", "sofa", "pottedplant", "bed",
    "diningtable", "toilet", "tvmonitor", "laptop", "mouse", "remote", "keyboard", "cell phone",
    "microwave", "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase", "scissors",
    "teddy bear", "hair drier", "toothbrush"
]

#导入掩码，在指定的区域进行检测
#mask=cv2.imread("../Image/mask.png")

#创建排序器的实例
tracker=Sort(max_age=20,min_hits=3,iou_threshold=0.3)

#limits=[723,297,1145,297]
limits = [350,597,1145,597]

totalcount=[]


while True:
    success,img=cap.read()

    ###加入图示文件，每次循环都要重新导入，否则会影响显示效果
    #imgGraphics=cv2.imread("../Image/graphics.png",cv2.IMREAD_UNCHANGED)
    #img=cvzone.overlayPNG(img,imgGraphics, (0,0))

    ###指定区域检测
    #imgRegion=cv2.bitwise_and(img,mask)

    ###指定区域检测的话，要将img换成imgRegion
    results=model(img,stream=True)

    #创建一个空列表来存储检测到的车辆数量
    detections=np.empty((0,5))


    for r in results:
        boxes=r.boxes
        for box in boxes:
            x1,y1,x2,y2=box.xyxy[0]



            x1, y1, x2, y2=int(x1),int(y1),int(x2),int(y2)

            print(x1,y1,x2,y2)


            w,h=x2-x1,y2-y1



            ###置信度
            conf=math.ceil((box.conf[0]*100))/100
            #cvzone.putTextRect(img,f'{conf}',(max(0,x1),max(35,y1)))

            ###处理类别名称
            cls=int(box.cls[0])
            currentClass=classNames[cls]

            if currentClass=="motorbike" or currentClass=="truck" or currentClass =="bus" \
                    and conf > 0.3:
                # cvzone.putTextRect(img, f'{classNames[cls]}{conf}', (max(0, x1), max(35, y1)),
                #                    scale=0.3, thickness=1, offset=3)

                #cvzone.cornerRect(img, (x1, y1, w, h), l=9)

                currentArray=np.array([x1,y1,x2,y2,conf])
                detections=np.vstack((detections,currentArray))


    resultsTracker=tracker.update(detections)
    cv2.line(img, (limits[0], limits[1]), (limits[2], limits[3]), (0, 0, 255), 2)

    for result in resultsTracker:
        x1,y1,x2,y2,Id=result
        x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
        w = x2 - x1
        h = y2 - y1
        print(result)
        cvzone.cornerRect(img,(x1,y1,w,h),l=0,rt=2,colorR=(255,0,0))

        cvzone.putTextRect(img,f'ID:{int(Id)}',(max(0,x1),max(35,y1)),
                       scale=0.6,thickness=1,offset=3)

        cx,cy=x1+w//2,y1+h//2
        cv2.circle(img,(cx,cy),5,(255,0,255),cv2.FILLED)

        if limits[0]<cx<limits[2] and limits[1]-30<cy<limits[1]+30:
            if totalcount.count(Id) ==0:
                totalcount.append(Id)
                cv2.line(img, (limits[0], limits[1]), (limits[2], limits[3]), (0, 0, 255), 2)


    #cvzone.putTextRect(img,f'Total Count:{len(totalcount)}',(50,50))
    cv2.putText(img,str(len(totalcount)),(255,100),cv2.FONT_HERSHEY_PLAIN,5,(50,50,255),8)

    cv2.imshow("image",img)
    cv2.waitKey(1)
